import pandas as pd
import sys

# 给数据集打上标签功能，argv[1]是输入文件，argv[2]是数据集标签
# argv[3]和argv[4]是可选，如果有，argv[3]是IP，argv[4]是标签
# argv[5]是输出文件

# 给数据集打上标签
def add_label(data, label):
    data.loc[:,'class']=label
    return data

# 给数据集某些IP打上标签
def add_IP_label(data, ip, label):
    data['class'][(data['srcIP']==ip)|(data['dstIP']==ip)]=label
    return data


if __name__ == '__main__':
    data = pd.read_csv(sys.argv[1], low_memory=False, delimiter='\t')
    data = add_label(data, sys.argv[2])
    if(len(sys.argv)==6):
        data = add_IP_label(data, sys.argv[3], sys.argv[4])
    
    data.to_csv(sys.argv[-1], index=False, sep=',')
